#!/Users/mabo/opt/anaconda3/bin/python
# -*- coding: utf-8 -*-
#Seismic Velocity-Derived Geophysical Models, can add other models or replace the empirical function
import os
import sys
import re
import numpy as np 
import pandas as pd
import matplotlib.pyplot as plt
import segyio
from colored import fg, bg, attr

C_GREEN = fg("green")
C_RED = fg("red")
C_BLUE = fg("blue")
C_DEFAULT = attr("reset")

#------------------
def plot_Vp(seyfile,cdp_inc=100,depth_inc=20):
    with segyio.open(os.path.join(os.getcwd(),'segy/',seyfile), ignore_geometry=True) as segyfile:
        n_traces = segyfile.tracecount
        headers = segyio.tracefield
        headers = segyio.tracefield.keys
        df = pd.DataFrame(index=range(1, n_traces + 1), columns=headers.keys())
        for k, v in headers.items():
            df[k] = segyfile.attributes(v)[:]
        data = segyfile.trace.raw[:]
        dataT = data.T
        twt = segyfile.samples
        cdp = df['CDP']
        #length = cdp * 6.25 / 1000
        raw = pd.DataFrame(dataT,index=twt,columns = cdp)
        raw = raw.rename_axis(index="twt", columns=["cdp"])
        seafloor=pd.read_csv('seafloor.txt',sep='\s+',header=None,names=['cdp','depth'],index_col='cdp') #seafloor file made by gmt
        seafloor=seafloor.loc[cdp.min():cdp.max()]
        Vp = raw.T.where(raw.index>seafloor.loc[seafloor.index]).T
        #Vp = raw.loc[np.arange(twt.min(),twt.max(),depth_inc),np.arange(cdp.min(),cdp.max(),cdp_inc)] 
        Vp = (Vp/1000).round(3)
        #Vp = Vp[(Vp > 1.5) | (Vp < 8.1)] #only for the data in 1.5-8.1 km/s
        Vp_inc = Vp.loc[np.arange(twt.min(),twt.max(),depth_inc),np.arange(cdp.min(),cdp.max(),cdp_inc)]
        ys,xs =np.mgrid[twt.min():twt.max():depth_inc,cdp.min():cdp.max():cdp_inc]
    fig, ax = plt.subplots(
        1,
        1,
        sharex=False,
        sharey=False,
        figsize=(12,12/6.25*1.5)
        )
    ax.set_aspect(1/6.25*1.5) #VE=1.5
    plt.xlim(cdp.min(),cdp.max())
    plt.ylim(twt.max(),twt.min())
    ax.set_title('Vp',weight='bold')
    ax.set_xlabel('CDP')
    ax.set_ylabel('Depth(m)')
    ax = plt.gca()
    cs = ax.contour(xs, ys,Vp_inc.values, levels=14, linewidths=0.5, colors='k')
    ax.clabel(cs, inline=1, fontsize=8)
    #norm = plt.Normalize(twt.min(), twt.max())
    cbar = ax.contourf(xs, ys,Vp_inc.values,levels=14, cmap="GnBu")
    cax = fig.add_axes([ax.get_position().x1+0.01,ax.get_position().y0,0.01,ax.get_position().height])
    cbar=plt.colorbar(cbar, cax=cax) # Similar to fig.colorbar(im, cax = cax)
    cbar.set_label(label='km/sec')
    ax.text(3000,28000, 'VE=1.5', dict(size=15))
    fig.savefig("Vp.png", dpi=300)

def plot_Vs(seyfile,cdp_inc=100,depth_inc=20):
    with segyio.open(os.path.join(os.getcwd(),'segy/',seyfile), ignore_geometry=True) as segyfile:
        n_traces = segyfile.tracecount
        headers = segyio.tracefield
        headers = segyio.tracefield.keys
        df = pd.DataFrame(index=range(1, n_traces + 1), columns=headers.keys())
        for k, v in headers.items():
            df[k] = segyfile.attributes(v)[:]
        data = segyfile.trace.raw[:]
        dataT = data.T
        twt = segyfile.samples
        cdp = df['CDP']
        #length = cdp * 6.25 / 1000
        raw = pd.DataFrame(dataT,index=twt,columns = cdp)
        raw = raw.rename_axis(index="twt", columns=["cdp"])
        seafloor=pd.read_csv('seafloor.txt',sep='\s+',header=None,names=['cdp','depth'],index_col='cdp') #seafloor file made by gmt
        seafloor=seafloor.loc[cdp.min():cdp.max()]
        Vp = raw.T.where(raw.index>seafloor.loc[seafloor.index]).T
        #Vp = raw.loc[np.arange(twt.min(),twt.max(),depth_inc),np.arange(cdp.min(),cdp.max(),cdp_inc)] 
        Vp = (Vp/1000).round(3)
        #Vp = Vp[(Vp > 1.5) | (Vp < 8.1)] #only for the data in 1.5-8.1 km/s
        Vs = 0.7858-1.2344*Vp+0.7949*Vp**2-0.1238*Vp**3+0.0064*Vp**4
        Vs = Vs.round(3)
        Vs_inc = Vs.loc[np.arange(twt.min(),twt.max(),depth_inc),np.arange(cdp.min(),cdp.max(),cdp_inc)]
        ys,xs =np.mgrid[twt.min():twt.max():depth_inc,cdp.min():cdp.max():cdp_inc]
    fig, ax = plt.subplots(
        1,
        1,
        sharex=False,
        sharey=False,
        figsize=(12,12/6.25*1.5)
        )
    ax.set_aspect(1/6.25*1.5) #VE=1.5
    plt.xlim(cdp.min(),cdp.max())
    plt.ylim(twt.max(),twt.min())
    ax.set_title('Vs(${0.7858-1.2344*Vp+0.7949*Vp^2-0.1238*Vp^3+0.0064*Vp^4(Brocher,2005)}$)',weight='bold')
    ax.set_xlabel('CDP')
    ax.set_ylabel('Depth(m)')
    ax = plt.gca()
    cs = ax.contour(xs, ys,Vs_inc.values, levels=14, linewidths=0.5, colors='k')
    ax.clabel(cs, inline=1, fontsize=8)
    #norm = plt.Normalize(twt.min(), twt.max())
    cbar = ax.contourf(xs, ys,Vs_inc.values,levels=14, cmap="GnBu")
    cax = fig.add_axes([ax.get_position().x1+0.01,ax.get_position().y0,0.01,ax.get_position().height])
    cbar=plt.colorbar(cbar, cax=cax) # Similar to fig.colorbar(im, cax = cax)
    cbar.set_label(label='km/sec')
    ax.text(3000,28000, 'VE=1.5', dict(size=15))
    fig.savefig("Vs.png", dpi=300)

def plot_density(seyfile,cdp_inc=100,depth_inc=20):
    with segyio.open(os.path.join(os.getcwd(),'segy/',seyfile), ignore_geometry=True) as segyfile:
        n_traces = segyfile.tracecount
        headers = segyio.tracefield
        headers = segyio.tracefield.keys
        df = pd.DataFrame(index=range(1, n_traces + 1), columns=headers.keys())
        for k, v in headers.items():
            df[k] = segyfile.attributes(v)[:]
        data = segyfile.trace.raw[:]
        dataT = data.T
        twt = segyfile.samples
        cdp = df['CDP']
        #length = cdp * 6.25 / 1000
        raw = pd.DataFrame(dataT,index=twt,columns = cdp)
        raw = raw.rename_axis(index="twt", columns=["cdp"])
        seafloor=pd.read_csv('seafloor.txt',sep='\s+',header=None,names=['cdp','depth'],index_col='cdp') #seafloor file made by gmt
        seafloor=seafloor.loc[cdp.min():cdp.max()]
        Vp = raw.T.where(raw.index>seafloor.loc[seafloor.index]).T
        #Vp = raw.loc[np.arange(twt.min(),twt.max(),depth_inc),np.arange(cdp.min(),cdp.max(),cdp_inc)] 
        Vp = (Vp/1000).round(3)
        Vp = Vp[(Vp > 1.5) | (Vp < 8.1)] #only for the data in 1.5-8.1 km/s
        density = 1.6612*Vp-0.4721*Vp**2+0.0671*Vp**3-0.0043*Vp**4+0.000106*Vp**5
        density = density.round(3)
        density_inc = density.loc[np.arange(twt.min(),twt.max(),depth_inc),np.arange(cdp.min(),cdp.max(),cdp_inc)]
        ys,xs =np.mgrid[twt.min():twt.max():depth_inc,cdp.min():cdp.max():cdp_inc]
    fig, ax = plt.subplots(
        1,
        1,
        sharex=False,
        sharey=False,
        figsize=(12,12/6.25*1.5)
        )
    ax.set_aspect(1/6.25*1.5) #VE=1.5
    plt.xlim(cdp.min(),cdp.max())
    plt.ylim(twt.max(),twt.min())
    ax.set_title('Density(${1.6612Vp-0.4721Vp^2+0.0671Vp^3-0.0043Vp^4+0.000106Vp^5}(Brocher,2005)$)',weight='bold')
    ax.set_xlabel('CDP')
    ax.set_ylabel('Depth(m)')
    ax = plt.gca()
    cs = ax.contour(xs, ys,density_inc.values, levels=14, linewidths=0.5, colors='k')
    ax.clabel(cs, inline=1, fontsize=8)
    #norm = plt.Normalize(twt.min(), twt.max())
    cbar = ax.contourf(xs, ys,density_inc.values,levels=14, cmap="GnBu")
    cax = fig.add_axes([ax.get_position().x1+0.01,ax.get_position().y0,0.01,ax.get_position().height])
    cbar=plt.colorbar(cbar, cax=cax) # Similar to fig.colorbar(im, cax = cax)
    cbar.set_label(label='${g/cm^2}$')
    ax.text(3000,28000, 'VE=1.5', dict(size=15))
    fig.savefig("density.png", dpi=300)

def plot_poisson(seyfile,cdp_inc=100,depth_inc=20):
    with segyio.open(os.path.join(os.getcwd(),'segy/',seyfile), ignore_geometry=True) as segyfile:
        n_traces = segyfile.tracecount
        headers = segyio.tracefield
        headers = segyio.tracefield.keys
        df = pd.DataFrame(index=range(1, n_traces + 1), columns=headers.keys())
        for k, v in headers.items():
            df[k] = segyfile.attributes(v)[:]
        data = segyfile.trace.raw[:]
        dataT = data.T
        twt = segyfile.samples
        cdp = df['CDP']
        #length = cdp * 6.25 / 1000
        raw = pd.DataFrame(dataT,index=twt,columns = cdp)
        raw = raw.rename_axis(index="twt", columns=["cdp"])
        seafloor=pd.read_csv('seafloor.txt',sep='\s+',header=None,names=['cdp','depth'],index_col='cdp') #seafloor file made by gmt
        seafloor=seafloor.loc[cdp.min():cdp.max()]
        Vp = raw.T.where(raw.index>seafloor.loc[seafloor.index]).T
        #Vp = raw.loc[np.arange(twt.min(),twt.max(),depth_inc),np.arange(cdp.min(),cdp.max(),cdp_inc)] 
        Vp = (Vp/1000).round(3)
        Vp = Vp[(Vp > 1.5) | (Vp < 8.5)] #only for the data in 1.5-8.5 km/s
        poisson = 0.8835-0.315*Vp+0.0491*Vp**2-0.0024*Vp**3
        poisson = poisson.round(3)
        poisson_inc = poisson.loc[np.arange(twt.min(),twt.max(),depth_inc),np.arange(cdp.min(),cdp.max(),cdp_inc)]
        ys,xs =np.mgrid[twt.min():twt.max():depth_inc,cdp.min():cdp.max():cdp_inc]
    fig, ax = plt.subplots(
        1,
        1,
        sharex=False,
        sharey=False,
        figsize=(12,12/6.25*1.5)
        )
    ax.set_aspect(1/6.25*1.5) #VE=1.5
    plt.xlim(cdp.min(),cdp.max())
    plt.ylim(twt.max(),twt.min())
    ax.set_title('Poisson(${0.8835-0.315*Vp+0.0491*Vp^2-0.0024*Vp^3}(Brocher,2005)$)',weight='bold')
    ax.set_xlabel('CDP')
    ax.set_ylabel('Depth(m)')
    ax = plt.gca()
    cs = ax.contour(xs, ys,poisson_inc.values, levels=14, linewidths=0.5, colors='k')
    ax.clabel(cs, inline=1, fontsize=8)
    #norm = plt.Normalize(twt.min(), twt.max())
    cbar = ax.contourf(xs, ys,poisson_inc.values,levels=14, cmap="GnBu")
    cax = fig.add_axes([ax.get_position().x1+0.01,ax.get_position().y0,0.01,ax.get_position().height])
    cbar=plt.colorbar(cbar, cax=cax) # Similar to fig.colorbar(im, cax = cax)
    ax.text(3000,28000, 'VE=1.5', dict(size=15))
    fig.savefig("poisson.png", dpi=300)

def plot_rigidity(seyfile,cdp_inc=100,depth_inc=20):
    #Rigidity G=rho*vs*vs
    with segyio.open(os.path.join(os.getcwd(),'segy/',seyfile), ignore_geometry=True) as segyfile:
        n_traces = segyfile.tracecount
        headers = segyio.tracefield
        headers = segyio.tracefield.keys
        df = pd.DataFrame(index=range(1, n_traces + 1), columns=headers.keys())
        for k, v in headers.items():
            df[k] = segyfile.attributes(v)[:]
        data = segyfile.trace.raw[:]
        dataT = data.T
        twt = segyfile.samples
        cdp = df['CDP']
        #length = cdp * 6.25 / 1000
        raw = pd.DataFrame(dataT,index=twt,columns = cdp)
        raw = raw.rename_axis(index="twt", columns=["cdp"])
        seafloor=pd.read_csv('seafloor.txt',sep='\s+',header=None,names=['cdp','depth'],index_col='cdp') #seafloor file made by gmt
        seafloor=seafloor.loc[cdp.min():cdp.max()]
        Vp = raw.T.where(raw.index>seafloor.loc[seafloor.index]).T
        #Vp = raw.loc[np.arange(twt.min(),twt.max(),depth_inc),np.arange(cdp.min(),cdp.max(),cdp_inc)] 
        Vp = (Vp/1000).round(3)
        #Vp = Vp[(Vp > 1.5) | (Vp < 8.5)] #only for the data in 1.5-8.1 km/s
        density = 1.6612*Vp-0.4721*Vp**2+0.0671*Vp**3-0.0043*Vp**4+0.000106*Vp**5
        Vs = 0.7858-1.2344*Vp+0.7949*Vp**2-0.1238*Vp**3+0.0064*Vp**4
        rigidity = density*Vs*Vs
        rigidity = rigidity.round(3)
        rigidity_inc = rigidity.loc[np.arange(twt.min(),twt.max(),depth_inc),np.arange(cdp.min(),cdp.max(),cdp_inc)]
        ys,xs =np.mgrid[twt.min():twt.max():depth_inc,cdp.min():cdp.max():cdp_inc]
    fig, ax = plt.subplots(
        1,
        1,
        sharex=False,
        sharey=False,
        figsize=(12,12/6.25*1.5)
        )
    ax.set_aspect(1/6.25*1.5) #VE=1.5
    plt.xlim(cdp.min(),cdp.max())
    plt.ylim(twt.max(),twt.min())
    ax.set_title('Rigidity(${density*Vs*Vs}(Sallares,2019)$)',weight='bold')
    ax.set_xlabel('CDP')
    ax.set_ylabel('Depth(m)')
    ax = plt.gca()
    cs = ax.contour(xs, ys,rigidity_inc.values, levels=14, linewidths=0.5, colors='k')
    ax.clabel(cs, inline=1, fontsize=8)
    #norm = plt.Normalize(twt.min(), twt.max())
    cbar = ax.contourf(xs, ys,rigidity_inc.values,levels=14, cmap="GnBu")
    cax = fig.add_axes([ax.get_position().x1+0.01,ax.get_position().y0,0.01,ax.get_position().height])
    cbar=plt.colorbar(cbar, cax=cax) # Similar to fig.colorbar(im, cax = cax)
    ax.text(3000,28000, 'VE=1.5', dict(size=15))
    fig.savefig("rigidity.png", dpi=300)

def plot_porosity(seyfile,cdp_inc=100,depth_inc=20):
    #Hyndman,1993
    with segyio.open(os.path.join(os.getcwd(),'segy/',seyfile), ignore_geometry=True) as segyfile:
        n_traces = segyfile.tracecount
        headers = segyio.tracefield
        headers = segyio.tracefield.keys
        df = pd.DataFrame(index=range(1, n_traces + 1), columns=headers.keys())
        for k, v in headers.items():
            df[k] = segyfile.attributes(v)[:]
        data = segyfile.trace.raw[:]
        dataT = data.T
        twt = segyfile.samples
        cdp = df['CDP']
        #length = cdp * 6.25 / 1000
        raw = pd.DataFrame(dataT,index=twt,columns = cdp)
        raw = raw.rename_axis(index="twt", columns=["cdp"])
        seafloor=pd.read_csv('seafloor.txt',sep='\s+',header=None,names=['cdp','depth'],index_col='cdp') #seafloor file made by gmt
        seafloor=seafloor.loc[cdp.min():cdp.max()]
        Vp = raw.T.where(raw.index>seafloor.loc[seafloor.index]).T
        #Vp = raw.loc[np.arange(twt.min(),twt.max(),depth_inc),np.arange(cdp.min(),cdp.max(),cdp_inc)] 
        Vp = (Vp/1000).round(3)
        #Vp = Vp[(Vp > 1.5) | (Vp < 8.5)] #only for the data in 1.5-8.1 km/s
        porosity = np.where(Vp.values<=3.9,1.941/Vp-0.2941,0.488888-0.074*Vp)
        porosity = porosity.round(3)*100
        porosity = pd.DataFrame(porosity,index=twt,columns=cdp)
        porosity_inc = porosity.loc[np.arange(twt.min(),twt.max(),depth_inc),np.arange(cdp.min(),cdp.max(),cdp_inc)]
        ys,xs =np.mgrid[twt.min():twt.max():depth_inc,cdp.min():cdp.max():cdp_inc]
    fig, ax = plt.subplots(
        1,
        1,
        sharex=False,
        sharey=False,
        figsize=(12,12/6.25*3)
        )
    ax.set_aspect(1/6.25*3)#VE=3
    plt.xlim(cdp.min(),cdp.max())
    plt.ylim(twt.max(),twt.min())
    ax.set_title('Porosity(if Vp<=3.9 do 1.941/Vp-0.2941 else 0.488888-0.074*Vp(Sallares,2019))',weight='bold')
    ax.set_xlabel('CDP')
    ax.set_ylabel('Depth(m)')
    ax = plt.gca()
    cs = ax.contour(xs, ys,porosity_inc.values, levels=14, linewidths=0.5, colors='k')
    ax.clabel(cs, inline=1, fontsize=8)
    #norm = plt.Normalize(twt.min(), twt.max())
    cbar = ax.contourf(xs, ys,porosity_inc.values,levels=14, cmap="GnBu")
    cax = fig.add_axes([ax.get_position().x1+0.01,ax.get_position().y0,0.01,ax.get_position().height])
    cbar=plt.colorbar(cbar, cax=cax) # Similar to fig.colorbar(im, cax = cax)
    cbar.set_label(label='%')
    ax.text(3000,28000, 'VE=3', dict(size=15))
    fig.savefig("porosity.png", dpi=300)

# ----------------------------------

def usage(argv):
    basename = argv[0]
    print("======================" + basename + "=======================")
    print("Seismic Velocity-Derived Geophysical Models")
    print(
        "["
        "Example(plot Vp)"
        " ]: " + C_BLUE + basename + C_RED + " -Vp" + C_DEFAULT
    )
    print(
        "["
        "Example(plot Vs)"
        " ]: " + C_BLUE + basename + C_RED + " -Vs" + C_DEFAULT
    )
    print(
        "["
        "Example(plot density)"
        " ]: " + C_BLUE + basename + C_RED + " -density" + C_DEFAULT
    )
    print(
        "["
        "Example(plot poisson)"
        " ]: " + C_BLUE + basename + C_RED + " -poisson" + C_DEFAULT
    )
    print(
        "["
        "Example(plot rigidity)"
        " ]: " + C_BLUE + basename + C_RED + " -rigidity" + C_DEFAULT
    )
    print(
        "["
        "Example(plot porosity)"
        " ]: " + C_BLUE + basename + C_RED + " -porosity" + C_DEFAULT
    )
    print("=======================================================")

def main(argv):
    seyfile = 'tmp_veld2.sgy' #change segy file here, segy file should put in segy/ file under current file
    if len(argv) < 2:
        usage(argv)
    else:
        if argv[1] == "-Vp":
            plot_Vp(seyfile,cdp_inc=100,depth_inc=20)
        elif argv[1] == "-Vs":
            plot_Vs(seyfile,cdp_inc=100,depth_inc=20)
        elif argv[1] == "-density":
            plot_density(seyfile,cdp_inc=100,depth_inc=20)
        elif argv[1] == "-poisson":
            plot_poisson(seyfile,cdp_inc=100,depth_inc=20)
        elif argv[1] == "-rigidity":
            plot_rigidity(seyfile,cdp_inc=100,depth_inc=20)
        elif argv[1] == "-porosity":
            plot_porosity(seyfile,cdp_inc=100,depth_inc=20)
        else:
            usage(argv)

if "__main__" == __name__:
    sys.exit(main(sys.argv))